package com.ruoyi.framework.interceptor;

import com.ruoyi.common.utils.TenantContext;
import org.apache.ibatis.executor.statement.StatementHandler;
import org.apache.ibatis.plugin.*;
import java.sql.Connection;
import java.util.Properties;

/**
 * MyBatis 租户拦截器，在SQL中自动注入tenant_id条件
 */
@Intercepts({
    @Signature(type = StatementHandler.class, method = "prepare", args = {Connection.class, Integer.class})
})
public class TenantInterceptor implements Interceptor {

    /**
     * 拦截SQL执行，动态注入tenant_id条件
     * @param invocation MyBatis拦截器调用对象
     * @return 执行结果
     * @throws Throwable 异常
     */
    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        StatementHandler handler = (StatementHandler) invocation.getTarget();
        String originalSql = handler.getBoundSql().getSql();
        Long tenantId = TenantContext.getTenantId();

        if (tenantId != null && needInjectTenantId(originalSql)) {
            String newSql = injectTenantIdCondition(originalSql, tenantId);
            // 反射设置新SQL
            java.lang.reflect.Field field = handler.getBoundSql().getClass().getDeclaredField("sql");
            field.setAccessible(true);
            field.set(handler.getBoundSql(), newSql);
        }
        return invocation.proceed();
    }

    /**
     * 判断是否需要注入租户条件（可根据表名、SQL类型等自定义过滤）
     * @param sql 原始SQL
     * @return 是否需要注入
     */
    private boolean needInjectTenantId(String sql) {
        String lowerSql = sql.trim().toLowerCase();
        return lowerSql.startsWith("select") || lowerSql.startsWith("update") || lowerSql.startsWith("delete");
    }

    /**
     * 在SQL中注入tenant_id条件
     * @param sql 原始SQL
     * @param tenantId 租户ID
     * @return 注入后的SQL
     */
    private String injectTenantIdCondition(String sql, Long tenantId) {
        String lowerSql = sql.toLowerCase();
        if (lowerSql.contains(" where ")) {
            return sql.replaceFirst("(?i)where", "WHERE tenant_id = " + tenantId + " AND ");
        } else if (lowerSql.contains(" order by ")) {
            int idx = lowerSql.indexOf(" order by ");
            return sql.substring(0, idx) + " WHERE tenant_id = " + tenantId + " " + sql.substring(idx);
        } else {
            return sql + " WHERE tenant_id = " + tenantId;
        }
    }

    @Override
    public Object plugin(Object target) {
        return Plugin.wrap(target, this);
    }

    @Override
    public void setProperties(Properties properties) {
        // 可扩展配置
    }
} 